package com.fanxl.flink.datastream.iterate;

import com.fanxl.flink.datastream.utils.NoOpIntMap;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.operators.ResourceSpec;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.IterativeStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoMapFunction;
import org.apache.flink.streaming.api.graph.StreamGraph;
import org.apache.flink.streaming.api.graph.StreamGraphGenerator;
import org.apache.flink.streaming.api.graph.StreamNode;
import org.junit.Test;

import static org.junit.Assert.*;

/**
 * @author Dell
 * @Title:
 * @Description: 请写注释类
 * @date 2022/7/11
 */
public class IterateITCase {


    @Test(expected = UnsupportedOperationException.class)
    public void testIncorrectParallelism() throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStream<Integer> source = env.fromElements(1, 10);

        IterativeStream<Integer> iter1 = source.iterate();
        SingleOutputStreamOperator<Integer> map1 = iter1.map(new NoOpIntMap());
        iter1.closeWith(map1).map(new NoOpIntMap());
    }


    /**
     * Test iteration job, check slot sharing group and co-location group.
     */
    @Test
    public void testIteration() {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStream<Integer> source = env.fromElements(1, 2, 3).name("source");
        IterativeStream<Integer> iteration = source.iterate(3000);
        iteration.name("iteration").setParallelism(2);
        DataStream<Integer> map = iteration.map(x -> {
            return x + 1;
        }).name("map").setParallelism(2);
        DataStream<Integer> filter = map.filter((x) -> {
            return false;
        }).name("filter").setParallelism(2);
        iteration.closeWith(filter).print();

        final ResourceSpec resources = ResourceSpec.newBuilder(1.0, 100).build();
        iteration.getTransformation().setResources(resources, resources);

        StreamGraph streamGraph = env.getStreamGraph();
        for (Tuple2<StreamNode, StreamNode> iterationPair : streamGraph.getIterationSourceSinkPairs()) {
            assertNotNull(iterationPair.f0.getCoLocationGroup());
            assertEquals(iterationPair.f0.getCoLocationGroup(), iterationPair.f1.getCoLocationGroup());

            assertEquals(StreamGraphGenerator.DEFAULT_SLOT_SHARING_GROUP, iterationPair.f0.getSlotSharingGroup());
            assertEquals(iterationPair.f0.getSlotSharingGroup(), iterationPair.f1.getSlotSharingGroup());

            final ResourceSpec sourceMinResources = iterationPair.f0.getMinResources();
            final ResourceSpec sinkMinResources = iterationPair.f1.getMinResources();
            final ResourceSpec iterationResources = sourceMinResources.merge(sinkMinResources);

        }
    }

    @Test
    public void testDoubleClosing() throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        // introduce dummy mapper to get to correct parallelism
        DataStream<Integer> source = env.fromElements(1, 10).map(new NoOpIntMap());

        IterativeStream<Integer> iter1 = source.iterate();

        iter1.closeWith(iter1.map(new NoOpIntMap()));
        iter1.closeWith(iter1.map(new NoOpIntMap()));
    }

    @Test
    public void testImmutabilityWithCoiteration() {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();


        DataStream<Integer> source = env.fromElements(1, 10).map(new NoOpIntMap()); // for rebalance

        IterativeStream<Integer> iter1 = source.iterate();
//         Calling withFeedbackType should create a new iteration
        IterativeStream.ConnectedIterativeStreams<Integer, String> iter2 = iter1.withFeedbackType(String.class);

        iter1.closeWith(iter1.map(new NoOpIntMap())).print("反馈流1");
        iter2.closeWith(iter2.map(noOpCoMap)).print("反馈流2");

        StreamGraph graph = env.getStreamGraph();

        assertEquals(2, graph.getIterationSourceSinkPairs().size());

        for (Tuple2<StreamNode, StreamNode> sourceSinkPair: graph.getIterationSourceSinkPairs()) {
            assertEquals(graph.getTargetVertex(sourceSinkPair.f0.getOutEdges().get(0)),
                    graph.getSourceVertex(sourceSinkPair.f1.getInEdges().get(0)));
        }
    }


    public static CoMapFunction<Integer, String, String> noOpCoMap =
            new CoMapFunction<Integer, String, String>() {

                public String map1(Integer value) throws Exception {
                    return value.toString();
                }

                public String map2(String value) throws Exception {
                    return value;
                }
            };
}
